/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph_optimizer/op_compiler/op_compiler_normal.h"
#include <unordered_set>
using namespace std;
namespace fe{

OpCompilerNormal::OpCompilerNormal(const std::string& compiler_name, const std::string& engine_name,
                                   OpStoreAdapterManagerPtr op_store_adapter_manager_ptr)
    :OpCompiler(compiler_name, engine_name, op_store_adapter_manager_ptr){}

OpCompilerNormal::~OpCompilerNormal() {}

Status OpCompilerNormal::Initialize(BufferFusionFunc &func) {
  // if graph optimizer has been initialized, return SUCCESS
  if (init_flag_) {
    FE_LOGW("OpCompiler has been initialized.");
    return SUCCESS;
  }

  init_flag_ = true;
  bufferFusionFunc_ = func;
  return SUCCESS;
}

void OpCompilerNormal::GetNodesNeedRePrcmpl(const Vistor<ge::NodePtr> &all_nodes,
                                            std::unordered_set<int64_t> &need_re_compile_scope_id,
                                            std::vector<ge::NodePtr> &nodes_be_compiled,
                                            std::vector<ge::NodePtr>& all_nodes_after_lx_fusion) {
  for (auto& node : all_nodes) {
    all_nodes_after_lx_fusion.emplace_back(node);
    bool need_re_compile = false;
    if (ge::AttrUtils::GetBool(node->GetOpDesc(), NEED_RE_PRECOMPILE, need_re_compile) && need_re_compile) {
      int64_t scope_id = 0;
      bool has_scope_id = GetFusionScopeAttr(node->GetOpDesc(), scope_id);
      if (has_scope_id && scope_id >= 0) {
        need_re_compile_scope_id.emplace(scope_id);
        /* fusion nodes will be emplaced into nodes_be_compiled in the following step.
         * Here only count all possible scope-ids. */
        FE_LOGD("Node %s need re-pe-compile and compile. We need to find all nodes with its scope id %ld.",
                node->GetName().c_str());
      } else {
        nodes_be_compiled.emplace_back(node);
        FE_LOGD("Single Node %s need re-pe-compile and compile.", node->GetName().c_str());
      }
    }
  }
}

Status OpCompilerNormal::GetFusionScope(ge::ComputeGraph& graph,
                                        ScopeNodeIdMap &fusion_nodes_map, std::vector<ge::NodePtr> &nodes_be_compiled,
                                        std::vector<ge::NodePtr>& all_nodes_after_lx_fusion,
                                        int64_t &minimum_negative_scope) {
  std::string graph_name = graph.GetName();
  auto all_nodes = graph.GetDirectNode();
  /* Find the minimum scope id in this graph. */
  for (auto& node : all_nodes) {
    (void)ge::AttrUtils::SetStr(node->GetOpDesc(), "_graph_name", graph_name);
    int64_t scope_id = 0;
    bool has_scope_id = GetFusionScopeAttr(node->GetOpDesc(), scope_id);
    if (has_scope_id && scope_id <= minimum_negative_scope) {
      minimum_negative_scope = scope_id - 1;
    }
  }

  /* Get all nodes which are need re-pre-compiled and store their scope id to
   * find all nodes with that scope id. */
  std::unordered_set<int64_t> need_re_compile_scope_id;
  GetNodesNeedRePrcmpl(all_nodes, need_re_compile_scope_id, nodes_be_compiled, all_nodes_after_lx_fusion);

  /* Find all nodes using special scope id. Every node uses that scope id
   * need to be re-compiled because they will be fused as one node. */
  for (auto& node : all_nodes) {
    int64_t scope_id = 0;
    bool has_scope_id = GetFusionScopeAttr(node->GetOpDesc(), scope_id);
    if (has_scope_id && scope_id >= 0) {
      if (need_re_compile_scope_id.count(scope_id) != 0) {
        FE_LOGD("node %s with scope id %ld needs to be re-compiled.", node->GetName().c_str(), scope_id);
        nodes_be_compiled.emplace_back(node);
      }
    }
  }

  if (GetScopeNodeMap(graph, nodes_be_compiled, minimum_negative_scope, fusion_nodes_map) != SUCCESS) {
    REPORT_FE_ERROR("[SubGraphOpt][Compile][GetFusScope] GetScopeNodeMap failed, graph [%s].", graph.GetName().c_str());
    return FAILED;
  }

  return SUCCESS;
}

Status OpCompilerNormal::ReCompileOpAfterLxFusion(ge::ComputeGraph& graph, CompileInfoParam &compile_info) {
  Status ret = PreCompileOp(graph);
  if (ret != SUCCESS) {
    REPORT_FE_ERROR("[SubGraphOpt][Compile][Lx] PreCompileOp failed after buffer fusion.");
    return ret;
  }

  /* scope_json_map is not need to be cleared. Put two rounds of compilation result together and
   * parse them once. */
  compile_info.minimum_negative_scope = -1;
  compile_info.fusion_nodes_map.clear();
  compile_info.buff_fus_to_del_nodes.clear();

  vector<ge::NodePtr> nodes_be_re_compiled;
  vector<ge::NodePtr> all_nodes;

  GetFusionScope(graph, compile_info.fusion_nodes_map, nodes_be_re_compiled,
                 all_nodes, compile_info.minimum_negative_scope);

  /* Using the same scope json map and we will parse json for all nodes including
   * those which are compiled before lx-fusion(and are not changed in lx-fusion) and those
   * which are changed in lx-fusion and need re-compiling. */
  ret = CompileOpOnly(graph, compile_info);
  if (ret != SUCCESS) {
    REPORT_INNER_ERROR(EM_INNER_ERROR,
        "[SubGraphOpt][Compile][Lx] Failed to compile after lx-fusion for graph %s with strategy %u and min scope %lu",
        graph.GetName().c_str(), compile_info.ignore_compile_strategy, compile_info.minimum_negative_scope);
    return ret;
  }

  ret = ParseJsonAndCompressOp(graph, compile_info.scope_json_map, all_nodes);
  if (ret != SUCCESS) {
    REPORT_FE_ERROR("[SubGraphOpt][Compile][Lx] Failed to parse json for graph [%s] after lx-fusion.",
                    graph.GetName().c_str());
    return ret;
  }
  return SUCCESS;
}

Status OpCompilerNormal::RunCompileProcess(ge::ComputeGraph& graph, const std::shared_ptr<GraphComm>& graph_comm_ptr,
                                           const std::shared_ptr<ScopeAllocator>& scope_allocator_ptr,
                                           std::vector<ge::NodePtr> &buff_fus_compile_failed_nodes,
                                           bool &need_post_process) {
  /* Some nodes needs to be re-pre-compiled after Ub fusion matching.
   * Because there format or data type. */
  Status ret;
  bool need_re_precompile_graph = false;
  (void)ge::AttrUtils::GetBool(graph, NEED_RE_PRECOMPILE, need_re_precompile_graph);
  if (need_re_precompile_graph) {
    ret = PreCompileOp(graph);
    if (ret != SUCCESS) {
      REPORT_FE_ERROR("[SubGraphOpt][Compile][Normal] PreCompileOp failed after buffer fusion for graph %s.",
                      graph.GetName().c_str());
      return ret;
    }
  }

  CompileInfoParam compile_info(buff_fus_compile_failed_nodes);
  vector<ge::NodePtr> nodes_be_compiled;
  vector<ge::NodePtr> buff_fus_rollback_nodes;

  ret = OpCompiler::GetFusionScope(graph, buff_fus_rollback_nodes, compile_info.fusion_nodes_map, nodes_be_compiled);
  if (ret != SUCCESS) {
    return SUCCESS;
  }

  /* Compile first to get a completely same graph as tuinng step which will lead us find
   * a optimization strategy. */
  FE_LOGI("Compile op and roll back failed fusion op first to ensure the graph will be successfully compiled.");
  FE_LOGI("Total size is %zu", nodes_be_compiled.size());
  ret = CompileOpOnly(graph, compile_info);
  if (ret != SUCCESS) {
    return ret;
  }

  FE_CHECK_NOTNULL(bufferFusionFunc_);
  ret = bufferFusionFunc_(graph, graph_comm_ptr, scope_allocator_ptr);
  /* Call Lxfusion to optimize the graph. */
  if (ret != SUCCESS && ret != tune::NO_FUSION_STRATEGY) {
    REPORT_FE_ERROR("[SubGraphOpt][Compile][Normal] BufferFusionProcess failed for graph %s result %u.",
                    graph.GetName().c_str(), ret);
    return FAILED;
  }

  // pre compile op which l1fusion or l2fusion changed
  if (ret != tune::NO_FUSION_STRATEGY) {
    FE_LOGI("Lx-fusion change the graph and we need re-precompile graph.");
    ret = ReCompileOpAfterLxFusion(graph, compile_info);
    if (ret != SUCCESS) {
      REPORT_FE_ERROR("[SubGraphOpt][Compile][RunCmplProc] Failed to re-compile op after lx fusion for graph [%s]",
                      graph.GetName().c_str());
      return FAILED;
    }
  } else {
    /* using the first compiling result only because no new nodes are created. */
    FE_LOGI("Lx-Fusion does not change the graph. We just use old compile result.");
    ret = ParseJsonAndCompressOp(graph, compile_info.scope_json_map, nodes_be_compiled);
    if (ret != SUCCESS) {
      REPORT_FE_ERROR("[SubGraphOpt][Compile][Normal] Failed to parse json or compress op for graph[%s].",
                      graph.GetName().c_str());
      return ret;
    }
  }
  FE_LOGI("Successfully compile op in normal mode.");
  need_post_process = true;
  return SUCCESS;
}

}
